{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 在CPA攻击中使用CW分析器"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "本篇将带你完成对软件AES实现的攻击，攻击例子是一个用C语言实现的AES算法，它和其他专用系统上的实现类似。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "SCOPETYPE = 'OPENADC'\n",
    "PLATFORM = 'CWLITEARM'\n",
    "CRYPTO_TARGET = 'TINYAES128C'\n",
    "num_traces = 50\n",
    "CHECK_CORR = False"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%bash -s \"$PLATFORM\" \"$CRYPTO_TARGET\"\n",
    "cd ../hardware/victims/firmware/simpleserial-aes\n",
    "make PLATFORM=$1 CRYPTO_TARGET=$2"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 捕获能轨"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 设置"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "下面使用了一些额外的帮助脚本减少编码量，如果使用的是XMEGA或者STM（CWLITEARM）板的话，下面应会运行正确。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%run \"Helper_Scripts/Setup_Generic.ipynb\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fw_path = '../hardware/victims/firmware/simpleserial-aes/simpleserial-aes-{}.hex'.format(PLATFORM)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "cw.program_target(scope, prog, fw_path)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "在捕获之前应该新创一个项目指定，之后的分析器需要用到它。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "project = cw.create_project(\"projects/Tutorial_B5\", overwrite = True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 捕获能轨"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "下面就是捕获的循环过程，每次加载一些新的明文，指定示波器，向目标发送密钥和明文，并且记录轨迹到项目中（一个包含trace_data, textin, textout, key的元组）。同时在我们绘图之前，还可以实时的查看数据。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%run \"Helper_Scripts/plot.ipynb\"\n",
    "plot = real_time_plot(plot_len=3000)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from tqdm import tnrange\n",
    "ktp = cw.ktp.Basic()\n",
    "\n",
    "for i in tnrange(num_traces, desc='Capturing traces'):\n",
    "    key, text = ktp.next()  # manual creation of a key, text pair can be substituted here\n",
    "    trace = cw.capture_trace(scope, target, text, key)\n",
    "    if trace is None:\n",
    "        continue\n",
    "    project.traces.append(trace)\n",
    "    plot.send(trace)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "然后就是绘图了"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%matplotlib notebook\n",
    "import matplotlib.pylab as plt\n",
    "plt.plot(project.waves[0])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "我们也可以其他库进行绘图，虽然matplotlib比较出名，但之后你会发现holoviews库很好用，这个库在处理大数据集或者是多条轨迹绘制上更优秀。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import holoviews as hv\n",
    "hv.extension('bokeh')\n",
    "hv.Curve(project.waves[0]).opts(width=600, height=600)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "最后，我们将数据存起来。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "project.save()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "CW到此工作结束，我们应该关闭其连接。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# cleanup the connection to the target and scope\n",
    "scope.dis()\n",
    "target.dis()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 分析"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import chipwhisperer as cw\n",
    "import chipwhisperer.analyzer as cwa\n",
    "proj = cw.open_project(\"projects/Tutorial_B5\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "现在我们拥有了能量轨迹，让我们来实现攻击吧！首先我们将泄漏模型设置为S盒的输出（在下一篇会介绍更多的泄漏模型）。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "leak_model = cwa.leakage_models.sbox_output"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "打印`cwa.leakage_models`，我们查看所有可用的泄漏模型。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(cwa.leakage_models)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "然后开始攻击"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "attack = cwa.cpa(proj, leak_model)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "打印attack对象我们可以看到攻击参数。你可以看到我们是用了所有的轨迹和所有的点，并且将攻击16字节密钥中每一位密钥。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(attack)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "最后，运行攻击。下面的代码是显示其原始输出，之后会使用其他的方法得到处理后的输出内容。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "results = attack.run()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "打印结果："
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(results)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "你会发现KGuess这一列和正确的密钥一致。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(results.find_key())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for key in proj.keys[0]:\n",
    "    print(\"0x{:02X}\".format(key))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "调用`results.find_maximums()`我们可以获取全部信息，它将返回如下内容：\n",
    "\n",
    "```Python\n",
    "find_maxiums() ->\n",
    "    [subkey0_data, subkey1_data, subkey2_data, ...]\n",
    "    \n",
    "subkey0_data ->\n",
    "    [guess0, guess1, guess2, ...]\n",
    "    \n",
    "guess0 ->\n",
    "    (key_guess, location_of_max, correlation)\n",
    "```\n",
    "\n",
    "举例，如果你想打印第四个子密钥的第三个最优猜测值的相关性。\n",
    "\n",
    "```python\n",
    "print(attack_results.find_maximums()[4][3][2])\n",
    "```\n",
    "\n",
    "注意，最大值的点坐标一般不会被计算或追逐，所以它将返回0，使用pandas库可以让我们在一个数据框中查看它们。同时我们必须对其进行转置才能获得期望的数据。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import pandas as pd\n",
    "stat_data = results.find_maximums()\n",
    "df = pd.DataFrame(stat_data).transpose()\n",
    "print(df.head())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "通常我们使用`.style`来进一步自定义，我们将使用链式函数调用（ROP），比如，我们可以去掉多余的0并且清理数据，由于我们知道正确的密钥，所以我们甚至可以用不懂的颜色来打印密钥。\n",
    "\n",
    "你可以在pandas库中使用大量的格式化方法，查看https://pandas.pydata.org/pandas-docs/stable/style.html了解详情。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "key = proj.keys[0]\n",
    "def format_stat(stat):\n",
    "    return str(\"{:02X}<br>{:.3f}\".format(stat[0], stat[2]))\n",
    "\n",
    "def color_corr_key(row):\n",
    "    global key\n",
    "    ret = [\"\"] * 16\n",
    "    for i,bnum in enumerate(row):\n",
    "        if bnum[0] == key[i]:\n",
    "            ret[i] = \"color: red\"\n",
    "        else:\n",
    "            ret[i] = \"\"\n",
    "    return ret\n",
    "\n",
    "df.head().style.format(format_stat).apply(color_corr_key, axis=1)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "你应该可以看到表格首行是红色数据（说明它是正确密钥），恭喜你已经完成了针对AES的CPA攻击。\n",
    "\n",
    "下面我们将看看如何使用分析器的其他功能来改进攻击过程，以及分析我们所拥有的数据。"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 间隔报告"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "当我们调用`attack.run()`时，我们是处理完所有的轨迹后才得到输出，这样对于比较短的攻击可行，但是对于较长的攻击，我们可以考虑在攻击中途获取一些反馈信息。可以通过给`attack.run()`传递一个回调函数来实现。这个函数将会间隔的被调用（默认是25轮轨迹，可以传入第二个参数进行控制）。\n",
    "\n",
    "让我们设置每个隔10轮轨迹便调用，我们只需要将现有代码放到回调函数真，并且增加`clear_output`清理输出并且使用`display`来实时显示。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from IPython.display import clear_output\n",
    "import numpy as np\n",
    "        \n",
    "def stats_callback():\n",
    "    results = attack.results\n",
    "    results.set_known_key(key)\n",
    "    stat_data = results.find_maximums()\n",
    "    df = pd.DataFrame(stat_data).transpose()\n",
    "    clear_output(wait=True)\n",
    "    display(df.head().style.format(format_stat).apply(color_corr_key,axis=1))\n",
    "    \n",
    "results = attack.run(stats_callback, 10)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "在jupyter中有默认的回调函数"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import chipwhisperer as cw\n",
    "cb = cwa.get_jupyter_callback(attack)\n",
    "results = attack.run(cb, 10)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "这便是我们每10轮轨迹轨迹边报告的结果，根据攻击以及你想学习的内容，你可以使用更高/低的值，一般来说频率越低会越快，但是更频繁的报告可以使你提前结束攻击。同时频繁报告也可以增加绘图数据的分辨率"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 绘制数据"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "分析器中还包含一个绘图模块用来帮助解释数据，它们每次作用于一个子键并返回可以用来绘制散景图（也可以选择其他模型）的数据。让我们从获取执行计算的类开始："
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### 输出 Vs. 时间"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plot_data = cwa.analyzer_plots(results)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "首先我们来查看输出/时间的模型，它将绘制出猜测值随着时间的相关性。这对于我们找到攻击操作的确切位置很有效，与之前的教程一样，我们使用散景图来绘制数据。\n",
    "\n",
    "对于`get_plot_data(bnum)`方法，它将返回一个列表`[xrange, correct_key, incorrect_key_data]`对于第`bnum`位，对于`incorrect_key_data`它返回两组错误数据，因为其中一组是正确猜测数据下方的的数据，另一组则为正确猜测数据上方的数据。\n",
    "\n",
    "我们会得到很多点，照常绘制即可，但是最后会抽取一些输出："
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def byte_to_color(idx):\n",
    "    return hv.Palette.colormaps['Category20'](idx/16.0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "import holoviews as hv\n",
    "from holoviews.operation.datashader import datashade, shade, dynspread, rasterize\n",
    "from holoviews.operation import decimate\n",
    "import pandas as pd, numpy as np\n",
    "\n",
    "a = []\n",
    "b = []\n",
    "hv.extension('bokeh')\n",
    "for i in range(0, 16):\n",
    "    data = plot_data.output_vs_time(i)\n",
    "    a.append(np.array(data[1]))\n",
    "    b.append(np.array(data[2]))\n",
    "    b.append(np.array(data[3]))\n",
    "    \n",
    "pda = pd.DataFrame(a).transpose().rename(str, axis='columns')\n",
    "pdb = pd.DataFrame(b).transpose().rename(str, axis='columns')\n",
    "curve = hv.Curve(pdb['0'], \"Sample\").options(color='black')\n",
    "for i in range(1, 16):\n",
    "    curve *= hv.Curve(pdb[str(i)]).options(color='black')\n",
    "    \n",
    "for i in range(0, 16):\n",
    "    curve *= hv.Curve(pda[str(i)]).options(color=byte_to_color(i))\n",
    "decimate(curve.opts(width=900, height=600))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "你会在图中看到一些尖峰，最大的一般就是S盒查找实际发送的地方（较小的一般是移动S盒数据的其他AES操作）。通常我们讨论的都是绝对值，所以会看到那里是负值。"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "这些信息在很多方面都有用，比如，你可以看到构成S盒查找的16个尖峰知识总轨迹的一小部分，如果我们需要重新运行攻击，我们可以捕获更少的样本以增加分析速度。"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### PGE vs. 轨迹"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "接下来我们将看到部分猜测熵（partial guessing entropy）/轨迹数量的模型。如前面所说，PGE是正确的子键在我们在我们猜测表中距离顶部的距离，比如，如果有7个猜测值的相关性高于实际子键，则该子键的PGE为7。\n",
    "\n",
    "该图对于查看攻击AES具体需要多少条轨迹很有效。但是记住绘图的分辨率由报告间隔决定（另外注意，必须在回调函数中调用`attack_results.find_maximums()`），在我们的例子中报告间隔为10，所以分辨率便为10条轨迹。\n",
    "\n",
    "方法与上图相似，因为它以`bnum`为参数并且返回`[xrange, PGE]`。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ret = plot_data.pge_vs_trace(0)\n",
    "curve = hv.Curve((ret[0],ret[1]), \"Traces Used in Calculation\", \"Partial Guessing Entrop of Byte\")\n",
    "for bnum in range(1, 16):\n",
    "    ret = plot_data.pge_vs_trace(bnum)\n",
    "    curve *= hv.Curve((ret[0],ret[1])).opts(color=byte_to_color(bnum))\n",
    "curve.opts(width=900, height=600)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "你可以看见许多行以高值开始然后迅速下降，你会发现其实攻击不需要使用所有的轨迹数据。\n",
    "\n",
    "尽管我们可以使用更少的轨迹攻击AES，但在实际攻击中，我们并不知道正确的密钥，所以我们不能使用PGE只能使用相关性来判断何时攻击成功。下面我们将了解如何使用更少的轨迹进行分析。"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### 相关性 vs. 轨迹"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "最后我们来看看相关性/轨迹数的关系图，和PGE/轨迹数一样，我们也是用10条轨迹的分辨率，这个方法将会返回一个`[xrange, [data_for_kguess]]`，所以我们需要绘制每个子键的每个猜测值图像，和上面一样，我们将为正确的子键绘制彩色其他的使用黑色。"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "你将会看到所有的猜测值最开始有一个很大的相关性，但是除了正确的之外其他值降迅速下降，思考一下，如果你不知道密钥，那么你在什么时候可以确定具有最高相关性的猜测实际上是正确的子密钥呢？\n",
    "\n",
    "让我们继续绘制正确的猜测和下一个最佳猜测的相关性："
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "a = []\n",
    "b = []\n",
    "for bnum in range(0, 16):\n",
    "    data = plot_data.corr_vs_trace(bnum)\n",
    "    best = [0] * len(data[1][0])\n",
    "    for i in range(255):\n",
    "        if i == key[bnum]:\n",
    "            a.append(np.array(data[1][i]))\n",
    "        else:\n",
    "            if max(best) < max(data[1][i]): best = data[1][i]\n",
    "    b.append(np.array(data[1][i]))\n",
    "\n",
    "pda = pd.DataFrame(a).transpose().rename(str, axis='columns')\n",
    "pdb = pd.DataFrame(b).transpose().rename(str, axis='columns')\n",
    "curve = hv.Curve(pdb['0'].tolist(), \"Iteration Number\", \"Max Correlation\").options(color='black')\n",
    "for i in range(1,len(pdb.columns)):\n",
    "    curve *= hv.Curve(pdb[str(i)]).options(color='black')\n",
    "    \n",
    "for i in range(len(pda.columns)):\n",
    "    curve *= hv.Curve(pda[str(i)]).options(color=byte_to_color(i))\n",
    "            \n",
    "curve.opts(width=900, height=600)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "和第一幅图一样，你可以看到彩色保存高位，而黑色线在下降，所以何时确定你已经攻击完了所有子键？以及它是否比所有PGE线都到达0时要高呢？"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 结论"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "你已经完成了一次成功的CPA攻击，并且了解了一些使用分析器改进攻击的功能。\n",
    "\n",
    "你可以转到更高级的教程，特别是像你展示实际攻击中的工作原理（手动CPA攻击教程），本教程还为ARM目标是用了`tiny-AES128-C`，它使用了与`XMEGA`目标相同的操作，之后的操作将会在更典型的AES上执行此攻击。"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Tests"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "key = list(proj.keys[0])\n",
    "recv_key = [kguess[0][0] for kguess in results.find_maximums()]\n",
    "assert (key == recv_key), \"Failed to recover encryption key\\nGot: {}\\nExpected: {}\".format(recv_key, key)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "assert (results.pge == [0]*16), \"PGE for some bytes not zero: {}\".format(results.pge)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "if CHECK_CORR:\n",
    "    max_corrs = [kguess[0][2] for kguess in results.find_maximums()]\n",
    "    assert (np.all([corr > 0.75 for corr in max_corrs])), \"Low correlation in attack (corr <= 0.75): {}\".format(max_corrs)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
